package com.lmk.dsa.algorithm.graph.miniTree;

import java.util.ArrayList;
import java.util.List;
import com.lmk.dsa.struct.commons.Edge;
import com.lmk.dsa.struct.commons.KvNode;
import com.lmk.dsa.struct.graph.Graph;
import com.lmk.dsa.struct.heap.MinIndexHeap;

/**
 * Prim最小生成树算法
 * @author LaoMake
 * @since 1.0
 */
public class Prim {

    /** 图的引用 */
    private Graph graph;

    /** 顶点的总数 */
    private int n;

    /** 访问状态数组 */
    private boolean[] marked;

    /** 标记是否索引节点是否已找到最短边 */
    private Edge[] edgeTo;

    /** 用于存储边的最小堆 */
    private MinIndexHeap heap;

    /** 最小生成树的边 */
    private List<Edge> edgeList;

    /** 最小生成树的总权值 */
    private int miniWeight;

    public Prim(Graph graph) {
        this.graph = graph;
        this.miniWeight = 0;
        this.n = graph.vertexSize();
        this.heap = new MinIndexHeap(graph.edgeSize());
        this.marked = new boolean[n];
        this.edgeTo = new Edge[n];
        this.edgeList = new ArrayList<>(n);
    }

    public int getMiniWeight() {
        return miniWeight;
    }

    public List<Edge> miniTree(){
        primVisit(0);
        KvNode<Integer, Edge> node;
        Edge edge;
        int index;
        while (heap.size() > 0){
            index = heap.poolIndex();
            edge = edgeTo[index];
            if (edge != null){
                edgeList.add(edge);
                primVisit(index);
            }
        }

        for (Edge e : edgeList){
            miniWeight += e.weight;
        }
        return edgeList;
    }

    private void primVisit(int v){
        List<Edge> edges;
        if(!marked[v]){
            marked[v] = true;
            edges = graph.adjoinEdge(v);
            int w;
            for (Edge edge : edges){
                w = edge.other(v);
                if(!marked[w]){
                    // 优化后的Prim算法，会提前丢弃无用边
                    if(edgeTo[w] == null){
                        heap.insert(w, edge.weight);
                        edgeTo[w] = edge;
                    }else if(edge.weight < edgeTo[w].weight){
                        edgeTo[w] = edge;
                        heap.change(w, edge.weight);
                    }
                }
            }
        }
    }
}
